#!/usr/bin/env python3
"""
conntrack-tools MCP Server

面向 conntrack-tools 软件包的 MCP 服务器，提供连接跟踪命令行操作的封装。

MCP tools 列表：
1. conntrack_list - 列出连接跟踪表
   用途：显示当前系统的连接跟踪表内容
   参数：table_name (str, 可选), output_format (str, 可选)
   返回：统一的命令执行结果 JSON

2. conntrack_stats - 显示连接跟踪统计信息
   用途：显示连接跟踪的统计计数器
   参数：无
   返回：统一的命令执行结果 JSON

3. conntrack_counters - 显示连接跟踪计数器
   用途：显示指定表的连接计数器
   参数：table_name (str, 可选)
   返回：统一的命令执行结果 JSON

4. conntrack_events - 监控连接跟踪事件
   用途：实时监控连接跟踪事件
   参数：event_mask (str, 可选), timeout (int, 可选)
   返回：统一的命令执行结果 JSON

5. conntrack_filter - 使用过滤条件列出连接
   用途：根据协议、地址族等条件过滤连接
   参数：protocol (str, 可选), family (str, 可选)
   返回：统一的命令执行结果 JSON

6. conntrack_zero_counters - 列出连接并重置计数器
   用途：列出连接跟踪表并重置计数器
   参数：无
   返回：统一的命令执行结果 JSON
"""

import subprocess
import shlex
from typing import Optional
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("Conntrack Tools MCP Server")

@mcp.tool()
def conntrack_list(table_name: Optional[str] = None, output_format: Optional[str] = None) -> dict:
    """
    列出连接跟踪表内容
    
    参数:
        table_name: 表名称，如 'dyning'，默认为全部表
        output_format: 输出格式，如 'xml'
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    cmd_parts = ["conntrack", "-L"]
    
    if table_name:
        cmd_parts.append(table_name)
    
    if output_format:
        cmd_parts.extend(["-o", output_format])
    
    command = " ".join(cmd_parts)
    
    try:
        result = subprocess.run(
            cmd_parts,
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def conntrack_stats() -> dict:
    """
    显示连接跟踪统计信息
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    command = "conntrack -S"
    
    try:
        result = subprocess.run(
            ["conntrack", "-S"],
            capture_output=True,
            text=True,
            timeout=10
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 10 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def conntrack_counters(table_name: Optional[str] = None) -> dict:
    """
    显示连接跟踪计数器
    
    参数:
        table_name: 表名称，默认为 'conntrack'
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    table = table_name or "conntrack"
    command = f"conntrack -C {table}"
    
    try:
        result = subprocess.run(
            ["conntrack", "-C", table],
            capture_output=True,
            text=True,
            timeout=10
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 10 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def conntrack_events(event_mask: Optional[str] = None, timeout: int = 10) -> dict:
    """
    监控连接跟踪事件
    
    参数:
        event_mask: 事件掩码，如 'NEW', 'DESTROY'
        timeout: 监控超时时间（秒），默认10秒
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    cmd_parts = ["timeout", str(timeout), "conntrack", "-E"]
    
    if event_mask:
        cmd_parts.extend(["-e", event_mask])
    
    command = " ".join(cmd_parts)
    
    try:
        result = subprocess.run(
            cmd_parts,
            capture_output=True,
            text=True,
            timeout=timeout + 5
        )
        
        return {
            "success": result.returncode in [0, 124],  # 124 is timeout exit code
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def conntrack_filter(protocol: Optional[str] = None, family: Optional[str] = None) -> dict:
    """
    使用过滤条件列出连接
    
    参数:
        protocol: 协议过滤，如 'tcp', 'udp', 'icmp'
        family: 地址族过滤，如 'ipv4', 'ipv6'
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    cmd_parts = ["conntrack", "-L"]
    
    if protocol:
        cmd_parts.extend(["-p", protocol])
    
    if family:
        cmd_parts.extend(["-f", family])
    
    command = " ".join(cmd_parts)
    
    try:
        result = subprocess.run(
            cmd_parts,
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def conntrack_zero_counters() -> dict:
    """
    列出连接跟踪表并重置计数器
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    command = "conntrack -L -z"
    
    try:
        result = subprocess.run(
            ["conntrack", "-L", "-z"],
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

if __name__ == "__main__":
    mcp.run()